package com.hunterai.comfyui;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;

import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.UUID;
import java.util.ArrayList;
import org.springframework.http.HttpMethod;
import java.util.Map;
// import java.util.HashMap;
import com.hunterai.util.RenameMultipartFile;
import org.apache.commons.io.FilenameUtils;
// import com.hunterai.comfyui.ComfyUIWorkflowManager;
import org.springframework.beans.factory.annotation.Autowired;


@Service
@Slf4j
public class ComfyUIService {

    @Value("${comfyui.api.url}")
    private String comfyUiApiUrl;
    
    // @Value("${comfyui.input.dir}")
    // private String inputDir;
    
    // @Value("${comfyui.output.dir}")
    // private String outputDir;
    
    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;
    private final ComfyUIWorkflowManager workflowManager;
    
    @Autowired
    public ComfyUIService(RestTemplate restTemplate, ObjectMapper objectMapper, ComfyUIWorkflowManager workflowManager) {
        this.restTemplate = restTemplate;
        this.objectMapper = objectMapper;
        this.workflowManager = workflowManager;
    }

    public String imageUpload(MultipartFile image) {
        // === 安全检查逻辑移到这里 ===
        String contentType = image.getContentType();
        if (contentType == null || !contentType.startsWith("image/")) {
            throw new IllegalArgumentException("只允许上传图片文件");
        }
        if (image.getSize() > 10 * 1024 * 1024) {
            throw new IllegalArgumentException("图片文件过大");
        }
        String ext = FilenameUtils.getExtension(image.getOriginalFilename());
        if (ext == null || (!ext.equalsIgnoreCase("jpg") && !ext.equalsIgnoreCase("jpeg")
                && !ext.equalsIgnoreCase("png") && !ext.equalsIgnoreCase("webp") && !ext.equalsIgnoreCase("bmp"))) {
            throw new IllegalArgumentException("不支持的图片格式");
        }
        // 生成随机文件名并上传
        String randomName = UUID.randomUUID().toString().replace("-", "");
        String newFileName = randomName + "." + ext;
        MultipartFile renamedImage = new RenameMultipartFile(image, newFileName);
    
        String uploadUrl = comfyUiApiUrl + "/upload/image";
        try {
            MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
            body.add("image", new MultipartInputStreamFileResource(renamedImage.getInputStream(), renamedImage.getOriginalFilename(), renamedImage.getSize()));
    
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.MULTIPART_FORM_DATA);
    
            HttpEntity<MultiValueMap<String, Object>> requestEntity = new HttpEntity<>(body, headers);
    
            ResponseEntity<String> response = restTemplate.postForEntity(uploadUrl, requestEntity, String.class);
    
            if (response.getStatusCode().is2xxSuccessful()) {
                String responseBody = response.getBody();
                ObjectMapper mapper = new ObjectMapper();
                JsonNode jsonNode = mapper.readTree(responseBody);
                String name = jsonNode.get("name").asText();
                return name;
            } else {
                throw new RuntimeException("上传图片到ComfyUI失败: " + response.getStatusCode());
            }
        } catch (Exception e) {
            throw new RuntimeException("上传图片到ComfyUI异常: " + e.getMessage(), e);
        }
    }
    
    /**
     * 执行ComfyUI工作流
     */
    public String executeWorkflow(JsonNode workflow) {
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        
        // 创建包含client_id的请求体
        ObjectNode requestBody = objectMapper.createObjectNode();
        
        // 生成随机的client_id，类似前端的实现
        String clientId = "java_" + UUID.randomUUID().toString().substring(0, 8);
        requestBody.put("client_id", clientId);
        requestBody.set("prompt", workflow);
        
        HttpEntity<String> requestEntity;
        try {
            requestEntity = new HttpEntity<>(objectMapper.writeValueAsString(requestBody), headers);
        } catch (Exception e) {
            throw new RuntimeException("Failed to serialize workflow", e);
        }
        
        // 发送请求到ComfyUI
        String result = restTemplate.postForObject(comfyUiApiUrl + "/prompt", requestEntity, String.class);
        
        return result;
    }

    public JsonNode generate(String workflowId, String imageName, Map<String, Object> params) throws Exception {
        // 1. 设置图片参数（如有）
        if (imageName != null) {
            params.put("image", imageName);
        }
    
        // 2. 应用参数到工作流
        JsonNode workflow = workflowManager.applyParams(workflowId, params);
    
        // 3. 执行工作流
        String result = executeWorkflow(workflow);
    
        // 4. 解析ComfyUI响应
        return objectMapper.readTree(result);
    }
    
    /**
     * 生成随机种子
     */
    public long generateRandomSeed() {
        return (long) (Math.random() * Long.MAX_VALUE);
    }

    /**
     * 检查任务状态
     */
    public TaskStatusResponse checkTaskStatus(String promptId) {
        try {
            String historyUrl = comfyUiApiUrl + "/history/" + promptId;  // 获取特定promptId的历史
            ResponseEntity<String> historyResponse = restTemplate.getForEntity(historyUrl, String.class);
            
            if (!historyResponse.getStatusCode().is2xxSuccessful()) {
                return TaskStatusResponse.error("Failed to fetch task history from ComfyUI");
            }
    
            JsonNode historyJson = objectMapper.readTree(historyResponse.getBody());
            
            // 如果返回空对象，说明任务还在处理中
            if (historyJson.isEmpty()) {
                return TaskStatusResponse.processing();
            }
    
            // 获取任务数据
            JsonNode taskData = historyJson.get(promptId);
            if (taskData.has("outputs")) {
                List<ImageInfo> images = extractImageInfo(taskData.get("outputs"), promptId);
                return TaskStatusResponse.success(images);
            }
            
            return TaskStatusResponse.processing();
    
        } catch (Exception e) {
            log.error("Error checking task status", e);
            return TaskStatusResponse.error(e.getMessage());
        }
    }

    /**
     * 获取图像数据
     */
    public ImageResponse getImage(String filename, String subfolder, String type) {
        try {
            String imageUrl = buildImageUrl(filename, subfolder, type);
            ResponseEntity<byte[]> response = restTemplate.exchange(
                    imageUrl,
                    HttpMethod.GET,
                    null,
                    byte[].class
            );

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                return new ImageResponse(response.getBody(), MediaType.IMAGE_PNG);
            } else {
                return ImageResponse.error("Failed to fetch image from ComfyUI");
            }
        } catch (Exception e) {
            log.error("Error fetching image", e);
            return ImageResponse.error(e.getMessage());
        }
    }

    private List<ImageInfo> extractImageInfo(JsonNode outputs, String promptId) {
        List<ImageInfo> images = new ArrayList<>();
        outputs.fields().forEachRemaining(outputEntry -> {
            JsonNode output = outputEntry.getValue();
            if (output.has("images")) {
                output.get("images").forEach(image -> {
                    String filename = image.get("filename").asText();
                    String subfolder = image.has("subfolder") ? image.get("subfolder").asText() : "";
                    String type = image.has("type") ? image.get("type").asText() : "output";
                    
                    images.add(new ImageInfo(
                        filename,
                        "image",
                        buildProxyImageUrl(promptId, filename, subfolder, type)
                    ));
                });
            }
        });
        return images;
    }

    private String buildImageUrl(String filename, String subfolder, String type) {
        StringBuilder url = new StringBuilder(comfyUiApiUrl)
            .append("/view?filename=").append(filename);
        
        if (!subfolder.isEmpty()) {
            url.append("&subfolder=").append(subfolder);
        }
        url.append("&type=").append(type);
        
        return url.toString();
    }

    private String buildProxyImageUrl(String promptId, String filename, String subfolder, String type) {
        return String.format("/cy/image?promptId=%s&filename=%s&subfolder=%s&type=%s",
                promptId, filename, subfolder, type);
    }    
}